import asyncio
import socket

clients = {}

class UDPServerProtocol:
    def connection_made(self, transport):
        self.transport = transport
        
    def datagram_received(self, data, addr):
        if addr not in clients:
            clients[addr] = self.transport
            print(f"Client {addr} connected.")
            
        # 转发数据给所有客户端
        for client_addr, client_transport in clients.items():
            if client_addr != addr:
                client_transport.sendto(data, client_addr)
                
    def connection_lost(self, exc):
        # 清理客户端信息
        for addr, transport in list(clients.items()):
            if transport == self.transport:
                print(f"Client {addr} disconnected.")
                del clients[addr]
                break

async def main():
    print("启动UDP中继服务器...")
    
    # 获取事件循环
    loop = asyncio.get_running_loop()
    
    # 创建UDP服务器
    transport, protocol = await loop.create_datagram_endpoint(
        lambda: UDPServerProtocol(),
        local_addr=('0.0.0.0', 9999))
    
    print("服务器正在运行，监听0.0.0.0:9999...")
    
    try:
        await asyncio.Future()  # 运行服务器直到被取消
    finally:
        transport.close()

if __name__ == "__main__":
    asyncio.run(main())